import cv2
import numpy as np
import torch
from torchvision import transforms
from model import FCN8s
import utils

def video_inference(model, weight_path, video_path, mean, std, size=(224, 224)):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # 加载模型权重
    checkpoint = torch.load(weight_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频: {video_path}")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        origin_h, origin_w = frame.shape[:2]
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame_resized = cv2.resize(frame_rgb, size)
        img_tensor = transform(frame_resized).unsqueeze(0).to(device)

        # 推理
        with torch.no_grad():
            output = model(img_tensor)
            pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()

        mask = (pred * 225).astype(np.uint8)
        mask = cv2.resize(mask, (origin_w, origin_h), interpolation=cv2.INTER_NEAREST)

        # 抠图
        result = cv2.bitwise_and(frame_rgb, frame_rgb, mask=mask)
        result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)

        # 显示
        cv2.imshow("Foreground", result_bgr)
        if cv2.waitKey(1) == ord('q'):
            break


    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    video_path = r"./demo.mp4"
    weight_path = r'../checkpoints/best_model.pth'
    mean, std = utils.calc_mean_std_func(img_dir="../Portrait-dataset-2000/dataset/training")

    # 初始化模型
    model = FCN8s(num_classes=2)

    video_inference(model, weight_path, video_path, mean, std)
